"""
10_gds_basics.py
================
Neo4j Graph Data Science (GDS) Fundamentals

This module covers:
1. GDS library overview
2. Graph projections
3. Memory estimation
4. Native vs Cypher projections
5. Graph catalog management
6. Running algorithms

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver

try:
    from graphdatascience import GraphDataScience
    GDS_AVAILABLE = True
except ImportError:
    GDS_AVAILABLE = False
    print("graphdatascience package not installed. Run: pip install graphdatascience")


@dataclass
class GraphProjection:
    """Information about a projected graph."""
    name: str
    node_count: int
    relationship_count: int
    memory_usage: str
    node_labels: List[str]
    relationship_types: List[str]


class GDSManager:
    """
    Manages Neo4j Graph Data Science operations.
    
    GDS Architecture:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │                    GDS WORKFLOW                                 │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │  1. PROJECT GRAPH                                               │
    │  ─────────────────                                              │
    │  Create in-memory graph from Neo4j                             │
    │                                                                 │
    │      Neo4j Database                                            │
    │           │                                                     │
    │           ▼                                                     │
    │      ┌─────────────────────────────────────┐                   │
    │      │     IN-MEMORY GRAPH PROJECTION      │                   │
    │      │                                     │                   │
    │      │  • Compressed adjacency list        │                   │
    │      │  • Node properties (optional)       │                   │
    │      │  • Relationship properties          │                   │
    │      │  • Optimized for traversal          │                   │
    │      └─────────────────────────────────────┘                   │
    │                                                                 │
    │  2. RUN ALGORITHM                                              │
    │  ────────────────                                               │
    │  Execute algorithm on projected graph                          │
    │                                                                 │
    │      • Centrality (PageRank, Betweenness, etc.)               │
    │      • Community (Louvain, Label Propagation, etc.)           │
    │      • Path finding (Dijkstra, A*, etc.)                       │
    │      • Similarity (Node Similarity, etc.)                      │
    │      • Embedding (Node2Vec, FastRP, etc.)                     │
    │                                                                 │
    │  3. RESULT MODES                                               │
    │  ──────────────                                                 │
    │      • stream: Return results as rows                          │
    │      • stats: Return summary statistics                        │
    │      • mutate: Add to projected graph                          │
    │      • write: Write back to Neo4j                              │
    │                                                                 │
    │  4. CLEANUP                                                     │
    │  ─────────                                                      │
    │  Drop projected graph to free memory                           │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, uri: str = None, username: str = None, password: str = None):
        """
        Initialize GDS connection.
        
        Args:
            uri: Neo4j URI (default: from environment)
            username: Username (default: from environment)
            password: Password (default: from environment)
        """
        if not GDS_AVAILABLE:
            raise ImportError("graphdatascience package required")
        
        uri = uri or os.getenv("NEO4J_URI", "bolt://localhost:7687")
        username = username or os.getenv("NEO4J_USERNAME", "neo4j")
        password = password or os.getenv("NEO4J_PASSWORD", "password123")
        
        self.gds = GraphDataScience(uri, auth=(username, password))
        self.driver = get_driver()
    
    @property
    def version(self) -> str:
        """Get GDS library version."""
        return self.gds.version()
    
    # =========================================================
    # GRAPH PROJECTIONS
    # =========================================================
    
    def create_native_projection(
        self,
        graph_name: str,
        node_labels: List[str],
        relationship_types: List[str],
        node_properties: List[str] = None,
        relationship_properties: List[str] = None,
        orientation: str = "NATURAL"
    ) -> GraphProjection:
        """
        Create a native graph projection.
        
        Native projections are faster but less flexible than Cypher projections.
        
        Orientations:
        - NATURAL: Keep original direction
        - REVERSE: Reverse all relationships
        - UNDIRECTED: Treat as undirected
        
        Args:
            graph_name: Name for the projection
            node_labels: Labels to include
            relationship_types: Relationship types to include
            node_properties: Node properties to load
            relationship_properties: Relationship properties to load
            orientation: Relationship orientation
            
        Returns:
            GraphProjection with metadata
        """
        # Drop if exists
        self.drop_graph(graph_name)
        
        # Build node projection
        node_projection = node_labels
        if node_properties:
            node_projection = {
                label: {"properties": node_properties}
                for label in node_labels
            }
        
        # Build relationship projection
        rel_projection = {
            rel_type: {"orientation": orientation}
            for rel_type in relationship_types
        }
        if relationship_properties:
            rel_projection = {
                rel_type: {
                    "orientation": orientation,
                    "properties": relationship_properties
                }
                for rel_type in relationship_types
            }
        
        # Create projection
        G, result = self.gds.graph.project(
            graph_name,
            node_projection,
            rel_projection
        )
        
        return GraphProjection(
            name=graph_name,
            node_count=result["nodeCount"],
            relationship_count=result["relationshipCount"],
            memory_usage=result.get("graphSize", "Unknown"),
            node_labels=node_labels,
            relationship_types=relationship_types
        )
    
    def create_cypher_projection(
        self,
        graph_name: str,
        node_query: str,
        relationship_query: str
    ) -> GraphProjection:
        """
        Create a Cypher-based graph projection.
        
        More flexible but slower than native projections.
        Use when you need filtering or computed properties.
        
        Node query must return: id, labels (optional), properties (optional)
        Relationship query must return: source, target, type (optional)
        
        Args:
            graph_name: Name for the projection
            node_query: Cypher query for nodes
            relationship_query: Cypher query for relationships
            
        Returns:
            GraphProjection with metadata
        """
        # Drop if exists
        self.drop_graph(graph_name)
        
        G, result = self.gds.graph.project.cypher(
            graph_name,
            node_query,
            relationship_query
        )
        
        return GraphProjection(
            name=graph_name,
            node_count=result["nodeCount"],
            relationship_count=result["relationshipCount"],
            memory_usage=result.get("graphSize", "Unknown"),
            node_labels=[],
            relationship_types=[]
        )
    
    def estimate_memory(
        self,
        node_labels: List[str],
        relationship_types: List[str]
    ) -> Dict[str, Any]:
        """
        Estimate memory required for projection.
        
        Use this before creating large projections to ensure
        you have sufficient memory.
        
        Args:
            node_labels: Labels to estimate
            relationship_types: Relationship types to estimate
            
        Returns:
            Memory estimation details
        """
        result = self.gds.graph.project.estimate(
            node_labels,
            relationship_types
        )
        
        return {
            "required_memory": result["requiredMemory"],
            "node_count": result["nodeCount"],
            "relationship_count": result["relationshipCount"]
        }
    
    # =========================================================
    # GRAPH CATALOG MANAGEMENT
    # =========================================================
    
    def list_graphs(self) -> List[Dict[str, Any]]:
        """List all projected graphs in the catalog."""
        result = self.gds.graph.list()
        return result.to_dict('records') if hasattr(result, 'to_dict') else []
    
    def graph_exists(self, graph_name: str) -> bool:
        """Check if a graph projection exists."""
        return self.gds.graph.exists(graph_name)["exists"]
    
    def get_graph(self, graph_name: str):
        """Get a reference to an existing graph projection."""
        if not self.graph_exists(graph_name):
            raise ValueError(f"Graph '{graph_name}' does not exist")
        return self.gds.graph.get(graph_name)
    
    def drop_graph(self, graph_name: str) -> bool:
        """Drop a graph projection from the catalog."""
        if self.graph_exists(graph_name):
            self.gds.graph.drop(graph_name)
            return True
        return False
    
    def drop_all_graphs(self):
        """Drop all graph projections."""
        graphs = self.list_graphs()
        for g in graphs:
            self.drop_graph(g["graphName"])
    
    # =========================================================
    # ALGORITHM EXECUTION PATTERNS
    # =========================================================
    
    def run_pagerank_stream(
        self,
        graph_name: str,
        damping_factor: float = 0.85,
        max_iterations: int = 20,
        tolerance: float = 0.0001
    ) -> List[Dict[str, Any]]:
        """
        Run PageRank in stream mode.
        
        Stream mode returns results without modifying
        the graph or database.
        
        Args:
            graph_name: Name of projected graph
            damping_factor: PageRank damping (0-1)
            max_iterations: Maximum iterations
            tolerance: Convergence tolerance
            
        Returns:
            List of nodes with PageRank scores
        """
        G = self.get_graph(graph_name)
        
        result = self.gds.pageRank.stream(
            G,
            dampingFactor=damping_factor,
            maxIterations=max_iterations,
            tolerance=tolerance
        )
        
        return result.to_dict('records')
    
    def run_pagerank_write(
        self,
        graph_name: str,
        write_property: str = "pagerank",
        damping_factor: float = 0.85
    ) -> Dict[str, Any]:
        """
        Run PageRank and write results to Neo4j.
        
        Write mode persists results to the database
        as node properties.
        
        Args:
            graph_name: Name of projected graph
            write_property: Property name for results
            damping_factor: PageRank damping
            
        Returns:
            Execution statistics
        """
        G = self.get_graph(graph_name)
        
        result = self.gds.pageRank.write(
            G,
            writeProperty=write_property,
            dampingFactor=damping_factor
        )
        
        return {
            "nodes_processed": result["nodePropertiesWritten"],
            "iterations": result["ranIterations"],
            "computation_millis": result["computeMillis"]
        }
    
    def run_louvain_stream(
        self,
        graph_name: str,
        include_intermediate: bool = False
    ) -> List[Dict[str, Any]]:
        """
        Run Louvain community detection in stream mode.
        
        Args:
            graph_name: Name of projected graph
            include_intermediate: Include intermediate communities
            
        Returns:
            List of nodes with community IDs
        """
        G = self.get_graph(graph_name)
        
        result = self.gds.louvain.stream(
            G,
            includeIntermediateCommunities=include_intermediate
        )
        
        return result.to_dict('records')
    
    def run_node_similarity_stream(
        self,
        graph_name: str,
        top_k: int = 10,
        similarity_cutoff: float = 0.0
    ) -> List[Dict[str, Any]]:
        """
        Run Node Similarity algorithm.
        
        Computes similarity based on shared neighbors.
        
        Args:
            graph_name: Name of projected graph
            top_k: Top K similar pairs per node
            similarity_cutoff: Minimum similarity
            
        Returns:
            List of similar node pairs
        """
        G = self.get_graph(graph_name)
        
        result = self.gds.nodeSimilarity.stream(
            G,
            topK=top_k,
            similarityCutoff=similarity_cutoff
        )
        
        return result.to_dict('records')
    
    # =========================================================
    # COMMON PATTERNS
    # =========================================================
    
    def analyze_graph_structure(
        self,
        graph_name: str
    ) -> Dict[str, Any]:
        """
        Analyze the structure of a projected graph.
        
        Runs multiple analyses to understand graph characteristics.
        """
        G = self.get_graph(graph_name)
        
        # Basic stats
        stats = {
            "name": graph_name,
            "node_count": G.node_count(),
            "relationship_count": G.relationship_count(),
            "density": G.relationship_count() / (G.node_count() ** 2) if G.node_count() > 0 else 0
        }
        
        # Degree statistics
        degree_result = self.gds.degree.stream(G)
        degrees = degree_result["score"].tolist()
        
        if degrees:
            stats["avg_degree"] = sum(degrees) / len(degrees)
            stats["max_degree"] = max(degrees)
            stats["min_degree"] = min(degrees)
        
        return stats
    
    def find_influential_nodes(
        self,
        graph_name: str,
        algorithm: str = "pagerank",
        top_n: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Find the most influential nodes in the graph.
        
        Args:
            graph_name: Name of projected graph
            algorithm: 'pagerank', 'betweenness', or 'degree'
            top_n: Number of top nodes to return
            
        Returns:
            Top influential nodes with scores
        """
        G = self.get_graph(graph_name)
        
        if algorithm == "pagerank":
            result = self.gds.pageRank.stream(G)
        elif algorithm == "betweenness":
            result = self.gds.betweenness.stream(G)
        elif algorithm == "degree":
            result = self.gds.degree.stream(G)
        else:
            raise ValueError(f"Unknown algorithm: {algorithm}")
        
        # Sort and get top N
        sorted_result = result.sort_values('score', ascending=False).head(top_n)
        
        return sorted_result.to_dict('records')


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("GDS Basics - Demonstration")
    print("=" * 60)
    
    if not GDS_AVAILABLE:
        print("❌ GDS library not available")
        exit(1)
    
    try:
        manager = GDSManager()
        
        print(f"\n📌 GDS Version: {manager.version}")
        print("-" * 40)
        
        # List existing graphs
        print("\n📌 Existing Graph Projections")
        print("-" * 40)
        graphs = manager.list_graphs()
        for g in graphs:
            print(f"   - {g['graphName']}")
        
        # Create a projection
        print("\n📌 Creating Graph Projection")
        print("-" * 40)
        
        try:
            projection = manager.create_native_projection(
                "healthcare_graph",
                ["Patient", "Disease", "Medication"],
                ["DIAGNOSED_WITH", "TAKES", "TREATS"],
                orientation="UNDIRECTED"
            )
            print(f"   Name: {projection.name}")
            print(f"   Nodes: {projection.node_count}")
            print(f"   Relationships: {projection.relationship_count}")
            
            # Run PageRank
            print("\n📌 Running PageRank")
            print("-" * 40)
            
            results = manager.run_pagerank_stream("healthcare_graph")
            for r in results[:5]:
                print(f"   Node {r['nodeId']}: {r['score']:.4f}")
            
            # Cleanup
            manager.drop_graph("healthcare_graph")
            print("\n   Graph dropped")
            
        except Exception as e:
            print(f"   Could not create projection: {e}")
        
        print("\n" + "=" * 60)
        print("✅ GDS basics demonstration complete!")
        print("=" * 60)
        
    except Exception as e:
        print(f"Error: {e}")
